from typing import List, Dict, Any, Optional
from langchain_milvus.vectorstores.milvus import Milvus as LangChainMilvus
from langchain_core.documents import Document


class FinalFixedMilvus(LangChainMilvus):
    """修正嵌套实体解析，不手动传递output_fields"""

    def _parse_document(self, data: Dict[str, Any]) -> Document:
        """修正后：直接从顶层data提取内容，适配实际数据结构"""
        print(f"\n=== 解析文档数据 ===")
        print(f"原始data（顶层字段）: {list(data.keys())}")  # 可验证字段是否存在

        # 核心修复：直接从顶层data获取page_content，而非嵌套的entity
        page_content = data.get("page_content", "")  # 从顶层提取
        print(f"提取的page_content（{self._text_field}）: {page_content[:50]}...")

        # 提取元数据：优先从顶层获取，适配实际结构
        metadata = {
            self._primary_field: data.get(self._primary_field),  # 顶层id
            "distance": data.get("distance")  # 若有distance也从顶层获取
        }



        print(f"生成的metadata: {metadata}")
        return Document(page_content=page_content, metadata=metadata)

    def _collection_search(
            self,
            embedding_or_text: Any,
            k: int,
            param: Dict[str, Any],
            expr: str = "",
            timeout: Optional[float] = None, **kwargs,
    ):
        # 完全不处理output_fields，交给父类自动处理
        print(f"执行搜索，k={k}，过滤条件={expr}")

        # 直接调用父类方法，不传递任何output_fields相关参数
        search_results = super()._collection_search(
            embedding_or_text=embedding_or_text,
            k=k,
            param=param,
            expr=expr,
            timeout=timeout,
            **kwargs
        )

        # 验证返回结构
        if search_results and len(search_results) > 0:
            first_hit_list = search_results[0]
            if first_hit_list and len(first_hit_list) > 0:
                first_hit = first_hit_list[0]
                print(f"第一条命中的完整结构: {first_hit}")
                entity_data = first_hit.get('entity', {})
                actual_fields = list(entity_data.keys()) if entity_data else []
                print(f"entity中实际的字段: {actual_fields}")
                if self._text_field not in actual_fields:
                    print(f"警告：entity中未找到 {self._text_field}，请检查数据插入是否正确")

        return search_results